from face_det_feature import FaceDetFeature
from db_operate import FaceDB

import cv2
import numpy as np
import json
from datetime import datetime, timezone
import os
import shutil
import sys

'''
# 获取人脸 
# 通过计算相似度，获取name
# 通过name查所有记录。
# 建立一个name文件夹，通过路经把所有文件复制到该路径。

python search_face.py ./test_img/test1.jpg
'''

def main():
    faceDB = FaceDB("face.db")
    faceID = {}
    rows = faceDB.select_all_faceid()
    for row in rows:
        faceID[row[0]] = {"name":row[1], "num":row[2], "feature":np.array(json.loads(row[3]))}
    
    faceDetFeature = FaceDetFeature()
    image_path = sys.argv[1]   # 测试图像  跳过脚本名(sys.argv[0]) "./test_img/test1.jpg" 
    img = cv2.imread(image_path)
    coords,features = faceDetFeature.get_info(img)
    # 找出ID，没有匹配就新增
    name = ""
    score = 0.0
    for key in faceID:
        score_temp = faceDetFeature.compare_features(faceID[key]["feature"], features[0])
        if score_temp>score:
            score = score_temp
            name = faceID[key]["name"]
    if score<0.40:
        print("没有匹配上的人脸：", score )
    else:
        print("匹配上人脸的name:", name)
        #name = "1"
        rows = faceDB.select_record_by_name(name)
        if os.path.exists(f"{name}") and os.path.isdir(f"{name}"):
            try:
                #os.rmdir(f"{name}")
                shutil.rmtree(f"{name}")
            except OSError as e:
                print(f"删除失败: {e}. 文件夹可能不为空。")
        os.mkdir(f"{name}")
        for index, row in enumerate(rows):
            img = cv2.imread(row[1])
            x1, y1, x2, y2 = map(int,json.loads(row[3]))
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img, f"Face: {name}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            cv2.imwrite(f'./{name}/{index}.jpg', img)
            #cv2.imshow(f'test', img)
            #if cv2.waitKey(0) or 0xFF == ord('q'):
            #    break


if __name__ == "__main__":
    main()


